import numpy as np


# 创建GridWorld环境
class GridWorld:
    areas_reward = np.array(
        [
            [0, 0, 0, 0, 0],
            [0, -10, -10, 0, 0],
            [0, 0, -10, 0, 0],
            [0, -10, 1, -10, 0],
            [0, -10, 0, 0, 0],
        ],
        dtype=float,
    )

    goal = np.array([3, 2])  # 目标状态

    def __init__(self, size):
        self.size = size
        self.state = np.array([0, 0])  # 初始状态

    @staticmethod
    def _transitions(_size):
        trans = np.zeros((_size, _size, 5, 2), dtype=int)
        for i in range(_size):
            for j in range(_size):
                trans[i][j][0] = np.array([i, j]) if i == 0 else np.array([i - 1, j])
                trans[i][j][1] = np.array([i, j]) if i == _size - 1 else np.array([i + 1, j])
                trans[i][j][2] = np.array([i, j]) if j == 0 else np.array([i, j - 1])
                trans[i][j][3] = np.array([i, j]) if j == _size - 1 else np.array([i, j + 1])
                trans[i][j][4] = np.array([i, j])
        return trans

    @property
    def transitions(self):
        return self._transitions(self.size)

    @property
    def state_space(self):
        return np.array([[y, x] for y in range(self.size) for x in range(self.size)])

    @property
    def action_space(self):
        return np.array(range(5))

    @staticmethod
    def _next_state(state, action, _size):
        i, j = state
        return GridWorld._transitions(_size)[i][j][action]

    def next_state(self, action):
        i, j = self.state
        return self.transitions[i][j][action]

    @staticmethod
    def _is_cross_border(state, action, _size):
        situations = [
            state[0] == 0 and action == 0,
            state[0] == _size - 1 and action == 1,
            state[1] == 0 and action == 2,
            state[1] == _size - 1 and action == 3,
        ]
        return True if any(situations) else False

    def is_cross_border(self, action):
        '''判断是否跨越边界'''
        return self._is_cross_border(self.state, action, self.size)

    @staticmethod
    def _reward(state, action, _size):
        i, j = GridWorld._next_state(state, action, _size)
        r1 = GridWorld.areas_reward[i][j]
        r2 = -1 if GridWorld._is_cross_border(state, action, _size) else 0
        return r1 + r2

    def reward(self, action):
        '''求取奖励'''
        return GridWorld._reward(self.state, action, self.size)

    def step(self, action):
        '''切换状态并求取奖励'''
        r = self.reward(action)
        self.state = self.next_state(action)
        done = True if all(self.state == self.goal) else False
        return self.state, r, done

    def reset(self, random=False):
        '''回到原点或者随机选取初始状态'''
        i = np.random.choice(self.size, 1)[0]
        j = np.random.choice(self.size, 1)[0]
        self.state = self.state_space[i * self.size + j] if random else np.array([0, 0])
